%% simulateForward.m
%
% Simulate the model forward as specified in initialize_sim_case
%--------------------------------------------------------------------------


%--------------------------------------------------------------------------
%% Initialization 
findStationary = 0;
g = g_stationary;
g_stacked = g_stationary(:);

AStore = AStoreSS;
interventionStore = interventionStoreSS;
cStore = cStoreSS;
rb_spot_ind = rb_spot_ind_stationary;

g_t_store = zeros(Nb,Na,Nz,Nr,length(rb_spot_ind_path));

meanc_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
means_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
meana_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
meanb_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
meancc_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
meannw_t_store = zeros(Nz,length(rb_spot_ind_path)+1);
rm_t_store = zeros(Nr*Nz,length(rb_spot_ind_path)+1);

    %get average consumption and assets (by income) before shocks
    income_shares = sum(squeeze(sum(sum(g(:,:,:,:)*da*db))),2);
    meanc = zeros(Nz,1);
    means = zeros(Nz,1);
    meana = zeros(Nz,1);
    meanb = zeros(Nz,1);
    meancc = zeros(Nz,1);
    for i = 1:Nr
        meanc = meanc + squeeze(sum(sum(cStore{rb_spot_ind, i}.*g(:,:,:,i)*da*db)))./income_shares;
        means = means + squeeze(sum(sum(sssStore{rb_spot_ind, i}.*g(:,:,:,i)*da*db)))./income_shares;
        meana = meana + squeeze(sum(sum(aaa.*g(:,:,:,i)*da*db)))./income_shares;
        meanb = meanb + squeeze(sum(sum(bbb.*g(:,:,:,i)*da*db)))./income_shares;
        meancc = meancc + squeeze(sum(sum(bbb.*(bbb < 0).*g(:,:,:,i)*da*db)))./income_shares;
    end
    meanc_t_store(:,1) = meanc;
    means_t_store(:,1) = means;
    meana_t_store(:,1) = meana;
    meanb_t_store(:,1) = meanb;
    meancc_t_store(:,1) = meancc;
    meannw_t_store(:,1) = (h - meana) + meanb;
    rm_t_store(:,1) = zeros(Nz*Nr,1);

time_path = [0, cumsum(Delta_path)];

if exist('switch_one_q', 'var') && switch_one_q == 1
    curr_slow_refi = slow_refi;
    slow_refi = 1-slow_refi;
end


%--------------------------------------------------------------------------
%% Simulate forward 
for t = 1:length(rb_spot_ind_path)    
    rb_spot_ind = rb_spot_ind_path(t);
    Delta = Delta_path(t);
    
    %Clear retirement matrix to allow for new spot rate
    if t == 1
        clear DD;
    elseif (rb_spot_ind_path(t) ~= rb_spot_ind_path(t-1))
        clear DD;
    end

    %Update policy functions for spot rate
    A = sparse([]);
    for rm_ind = 1:Nr
        A = [A; [sparse(M,M*(rm_ind-1)), AStore{rb_spot_ind,rm_ind}, sparse(M,M*(Nr-rm_ind))]];
    end

    %will always have g_stacked from previous iterations
    if ismember(distShockType, [1,2,3]) & ismember(t, distShockPeriods)
        %SC is shock intervention matrix
        g_stacked = transpose(SC)*g_stacked;
    end
    
    %Reset adjustment speed after 1 quarter
    if exist('switch_one_q', 'var') && switch_one_q == 1 && time_path(t+1) >= 0.25 && time_path(t) < 0.25
        slow_refi = curr_slow_refi;
    end
    
    solveKF;

    %store output
    g_t_store(:,:,:,:,t) = g;
    
    %get average consumption and assets (by income)
    income_shares = sum(squeeze(sum(sum(g(:,:,:,:)*da*db))),2);
    meanc = zeros(Nz,1);
    means = zeros(Nz,1);
    meana = zeros(Nz,1);
    meanb = zeros(Nz,1);
    meancc = zeros(Nz,1);
    for i = 1:Nr
        meanc = meanc + squeeze(sum(sum(cStore{rb_spot_ind, i}.*g(:,:,:,i)*da*db)))./income_shares;
        means = means + squeeze(sum(sum(sssStore{rb_spot_ind, i}.*g(:,:,:,i)*da*db)))./income_shares;
        meana = meana + squeeze(sum(sum(aaa.*g(:,:,:,i)*da*db)))./income_shares;
        meanb = meanb + squeeze(sum(sum(bbb.*g(:,:,:,i)*da*db)))./income_shares;
        meancc = meancc + squeeze(sum(sum(bbb.*(bbb < 0).*g(:,:,:,i)*da*db)))./income_shares;
    end
    meanc_t_store(:,t+1) = meanc;
    means_t_store(:,t+1) = means;
    meana_t_store(:,t+1) = meana;
    meanb_t_store(:,t+1) = meanb;
    meancc_t_store(:,t+1) = meancc;
    meannw_t_store(:,t+1) = (h - meana) + meanb;
        tmp=sum(sum(g*da*db));
    rm_t_store(:,t+1) = [squeeze(tmp(1,1,1,:))./income_shares(1);squeeze(tmp(1,1,2,:))./income_shares(2);squeeze(tmp(1,1,3,:))./income_shares(3)];

    display(['Iteration: ' num2str(t)]);
end
